import numpy as np
import matplotlib.pyplot as plt



# from mpl_toolkits.mplot3d import Axes3D

def mesh_single(mat, color):
    # 获取矩阵的形状
    # rows, cols = mat.shape
    if mat.shape[0]>80:
        coef = mat.shape[0] / 20
    else:
        coef = 1

    # 固定尺寸
    mat = mat[::int(coef), ::int(coef)]
    rows, cols = mat.shape

    mat = np.resize(mat, [rows, cols])

    # 创建网格 X 和 Y
    X, Y = np.meshgrid(np.arange(cols), np.arange(rows))

    # 创建图形

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # 绘制 3D 曲面图，使用矩阵 mat 作为 Z 值
    surf = ax.plot_surface(X, Y, mat, cmap=color, alpha=0.5)

    # 添加颜色条
    fig.colorbar(surf)

    # 显示图形
    plt.show()

def mesh_double(mat1, mat2, color1, color2, if_normalization, legend):
    # 获取矩阵的形状
    # rows, cols = mat.shape
    coef = int(np.ceil(mat1.shape[0] / 200))

    # 固定尺寸
    mat1 = mat1[::coef, ::coef]
    mat2 = mat2[::coef, ::coef]
    rows, cols = mat1.shape

    # 是否归一化
    if if_normalization:
        mat1 = mat1 / np.mean(mat1)
        mat2 = mat2 / np.mean(mat2)

    # 创建网格 X 和 Y
    X, Y = np.meshgrid(np.arange(cols), np.arange(rows))

    # 创建图形

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # 绘制 3D 曲面图，使用矩阵 mat 作为 Z 值
    ax.plot_surface(X, Y, mat1, cmap=color1, label=legend[0], alpha=0.7)
    ax.plot_surface(X, Y, mat2, cmap=color2, label=legend[1], alpha=0.7)

    # 使用代理艺术家创建图例
    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], color='red', lw=4, label=legend[0]),
        Line2D([0], [0], color='green', lw=4, label=legend[1])
    ]

    # 添加图例
    ax.legend(handles=legend_elements, loc='upper right')

    # 显示图形
    plt.show()
